# -*- coding: utf-8 -*-
# @Time    : 2024/9/8 20:13
# @Author  : Z.J. Zhang
# @Email   : zijingzhang@mail.ustc.edu.cn
# @File    : magnetic.py
# @Software: PyCharm
import abc
import typing

import matplotlib

matplotlib.use('tkagg')
import matplotlib.pyplot as plt
import numpy

import mot.field.base as base
import scipy.constants as C
from scipy.special import ellipk, ellipe
from myutils.test_utils import timer
import enum


class FieldNames(enum.Enum):
    Bx = 0
    By = 1
    Bz = 2
    B = 3


class BGenerator(base.FieldGenerator):
    field_names = FieldNames
    @abc.abstractmethod
    def B(self, x, y, z):
        pass

    @timer
    def cache(self, grid: base.RectGrid3D):
        B = self.B(grid.X, grid.Y, grid.Z, )
        grid.set_parameter_data({
            FieldNames.Bx.name: B[:, :, :, 0],
            FieldNames.By.name: B[:, :, :, 1],
            FieldNames.Bz.name: B[:, :, :, 2],
            FieldNames.B.name: numpy.linalg.norm(B, axis=3),
        })
        return grid


class CurrentCarryingLoop(BGenerator):
    """
    ref: 2011 Atom trapping in non-trivial geometries for micro-fabrication applications MVAngeleyn_Thesis_final.pdf
    page 136 Eq 6.8
    """

    def __init__(self, I, R, center: numpy.ndarray = numpy.array([0., 0., 0.]),
                 ez_of_coil_coord_sys: numpy.ndarray = numpy.array([0., 0., 1.]),
                 # ex_of_coil_coord_sys: numpy.ndarray = numpy.array([0., 0., 1.])
                 ):
        """

        :param I:
        :param R:
        :param center: 线圈的中心点
        :param ez_of_coil_coord_sys: 线圈的法向，根据右手螺旋得到
        """
        self.I = I
        self.R = R
        self.center = center
        self._ez_of_coil_coord_sys = ez_of_coil_coord_sys / numpy.linalg.norm(ez_of_coil_coord_sys)
        self._ex_of_coil_coord_sys = numpy.zeros(self._ez_of_coil_coord_sys.shape)
        self._ex_of_coil_coord_sys[self._ez_of_coil_coord_sys == 0] = 1.
        __where_is_not_0 = numpy.where(self._ez_of_coil_coord_sys != 0)
        self._ex_of_coil_coord_sys[__where_is_not_0[0][:-1]] = 1.
        self._ex_of_coil_coord_sys[__where_is_not_0[0][-1]] = -numpy.dot(self._ez_of_coil_coord_sys,
                                                                         self._ex_of_coil_coord_sys) / \
                                                              self._ez_of_coil_coord_sys[__where_is_not_0[0][-1]]
        self._ex_of_coil_coord_sys /= numpy.linalg.norm(self._ex_of_coil_coord_sys)

        self._ey_of_coil_coord_sys = numpy.cross(self._ez_of_coil_coord_sys, self._ex_of_coil_coord_sys)
        self._ey_of_coil_coord_sys /= numpy.linalg.norm(self._ey_of_coil_coord_sys)
        self._M_exeyez_of_coil_coord_sys = numpy.matrix([self._ex_of_coil_coord_sys,
                                                         self._ey_of_coil_coord_sys,
                                                         self._ez_of_coil_coord_sys])
        self._M_exeyez_of_coil_coord_sys_inv = self._M_exeyez_of_coil_coord_sys ** -1

        self._vec_B = numpy.vectorize(self._B,  # otypes=[float,float,float]
                                      signature="(),(),()-> (3)"
                                      )

    def B(self, x, y, z):
        return self._vec_B(x, y, z)

    def __alpha(self, r): return r / self.R

    def __beta(self, z): return z / self.R

    # def __Q(self,r,z):return (1+self.__alpha(r))**2 +self.__beta(z)**2
    def __Q(self, alpha, beta): return (1 + alpha) ** 2 + beta ** 2

    @staticmethod
    def __m(alpha, Q):
        return 4 * alpha / Q

    @staticmethod
    def __gamma(z, r): return z / r

    # @numpy.vectorize(signature="(),(),()-> (3)")
    def _B(self, x_wcs, y_wcs, z_wcs):
        # 得到Coil坐标系下的坐标
        x, y, z = ((numpy.array((x_wcs, y_wcs, z_wcs)) - self.center) * self._M_exeyez_of_coil_coord_sys_inv).A[0]

        r = (x ** 2 + y ** 2) ** 0.5
        er = numpy.array((x / r, y / r, 0))
        ez = numpy.array((0, 0, 1))
        alpha = self.__alpha(r)
        beta = self.__beta(z)

        Q = self.__Q(alpha, beta)
        m = self.__m(alpha, Q)
        gamma = self.__gamma(z, r)
        temp1 = Q - 4 * alpha
        Em = ellipe(m)
        Km = ellipk(m)
        B = (self.I * C.mu_0 / (2 * numpy.pi * self.R * Q ** 0.5) * (
                gamma * ((1 + alpha ** 2 + beta ** 2) / temp1 * Em - Km) * er
                + ((1 - alpha ** 2 - beta ** 2) / temp1 * Em + Km) * ez
        ) * self._M_exeyez_of_coil_coord_sys).A[0]

        return B
class BGenerators(BGenerator):
    def __init__(self, b_generators :typing.List[BGenerator]):
        self.b_generators = b_generators
    def B(self,x,y,z):
        res =0
        for bg in self.b_generators:
            res+= bg.B(x,y,z)
        return res

if __name__ == '__main__':
    plt.ion()
    xs = numpy.linspace(-10e-3, 10e-3, 20)
    ccl = CurrentCarryingLoop(200, 5e-3, numpy.array([0, 0, 5e-3]), numpy.array([1, 0, 0]))

    # plt.figure()
    # plt.plot(xs, ccl.B(xs, 0e-3, 0e-3) )

    X, Z = numpy.meshgrid(xs, xs, indexing='ij')
    _B = ccl.B(X, 0, Z, )
    plt.figure()
    plt.quiver(X, Z, _B[:, :, 0], _B[:, :, 2])
    plt.gca().set_aspect("equal")

    grid = base.default_grid

    ccl.cache(grid)
    grid_B = grid.get(X, 0, Z, ["Bx", "Bz"])

    plt.figure()
    plt.contourf(X, Z, grid.get(X, 0, Z, ['B'])[..., 0])
    plt.quiver(X, Z, grid_B[:, :, 0], grid_B[:, :, 1])
    plt.gca().set_aspect("equal")

    plt.figure()
    plt.contourf(X, Z, (grid_B[..., 0] ** 2 + grid_B[..., 1] ** 2) ** 0.5)
    plt.quiver(X, Z, grid_B[:, :, 0], grid_B[:, :, 1])
    plt.gca().set_aspect("equal")
